import sensor, image, time, math
from ulab import numpy as np
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA)
sensor.skip_frames(30)
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(False)
sensor.set_vflip(True)        #视角垂直翻转
sensor.set_hmirror(True)      # 视角水平翻转
clock = time.clock()
Ts = 1
# 状态空间矩阵定义
A = np.array([[1,0,0,0,Ts,0],[0,1,0,0,0,Ts],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]])
C = np.array([[1,0,0,0,0,0],[0,1,0,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]])
# 过程噪声协方差矩阵
Q_value = [1e-8 for _ in range(6)]
Q = np.diag(Q_value) #创建对角矩阵
# 观测噪声协方差矩阵
R_value = [1e-6 for _ in range(6)]
R = np.diag(R_value)
# 定义观测量Z
x = 0 #左顶点x坐标
y = 0 #左顶点y坐标
last_frame_x = x #上一帧左顶点x坐标
last_frame_y = y #上一帧左顶点y坐标
last_frame_location = [0 for _ in range(4)]
last_frame_rect = [0 for _ in range(4)]
box = [0 for _ in range(4)] #生成一个包含四个子列表的列表,存储x1,y1,x2,y2
tag_families = image.TAG36H11 #定义AprilTag家族，使用TAG36H11
w = 0 #矩形框宽度w
h = 0 #矩形框高度h
dx = 0 #左顶点x坐标移动速度
dy = 0 #左顶点y坐标移动速度
Z = np.array([x,y,w,h,dx,dy])
# 定义卡尔曼滤波函数变量
x_hat = np.array([80,60,30,30,2,2])
x_hat_minus = np.array([0,0,0,0,0,0])
p_value = [10 for _ in range(6)]
p = np.diag(p_value)
flag_action = 0
error_range = 7e-1
# 卡尔曼滤波函数
def Kalman_Filter(Z):
    global A,C,Q,R,x_hat,x_hat_minus,p
    # 预测部分
    x_hat_minus = np.dot(A,x_hat)
    p_minus = np.dot(A, np.dot(p, A.T)) + Q
    # 校正部分
    S = np.dot(np.dot(C, p_minus), C.T) + R
    # 选择一个小的正则化项
    regularization_term = 1e-4
    # 正则化 S 矩阵
    S_regularized = S + regularization_term * np.eye(S.shape[0])
    # 计算正则化后的 S 矩阵的逆
    S_inv = np.linalg.inv(S_regularized)
    # 计算卡尔曼增益 K
    K = np.dot(np.dot(p_minus, C.T), S_inv)
    x_hat = x_hat_minus + np.dot(K,(Z - np.dot(C,x_hat_minus)))
    p = np.dot((np.eye(6) - np.dot(K,C)),p_minus)
    return x_hat
def calculate_IoU(box1,box2):
    # 解析矩形框坐标
    x1, y1, x2, y2 = box1
    x3, y3, x4, y4 = box2
    #交集左上角点坐标
    Left_Peak_X = max(x1, x3)
    Left_Peak_Y = max(y1, y3)
    #交集右下角点坐标
    Right_Peak_X = min(x2, x4)
    Right_Peak_Y = min(y2, y4)
    if (Right_Peak_X - Left_Peak_X) > 0 and (Right_Peak_Y - Left_Peak_Y) > 0:
        #交集面积
        intersection_s = (Right_Peak_X - Left_Peak_X) * (Right_Peak_Y - Left_Peak_Y)
        #上一帧的矩形框面积s1
        s1 = (x2 - x1) * (y2 - y1)
        #当前帧的矩形框面积s2
        s2 = (x4 - x3) * (y4 - y3)
        #并集面积
        union_s = (s1 + s2 ) - intersection_s
        #交并比
        IoU = intersection_s / union_s
        return IoU
    else:
        return 0

while(True):
    clock.tick()
    img = sensor.snapshot()
    if flag_action == 0:
        for tag in img.find_apriltags():
            if tag:
                img.draw_rectangle(tag.rect(), color = (255, 0, 0))
                rect = tag.rect()
                box = [rect[0], rect[1], rect[0] + rect[2], rect[1] + rect[3]]
                last_frame_rect = rect
                last_frame_location = box
                x, y, w, h = rect[0], rect[1], rect[2], rect[3]
                x_hat = np.array([x,y,w,h,0,0])
                last_frame_x, last_frame_y = x, y
                flag_action = 1
                print('First tag detected and tracking started.')
    if flag_action == 1:
        flag_found_tag = False
        tags = img.find_apriltags(families=tag_families)  # 在图像中寻找AprilTag
        if tags:
            for tag in tags:
                img.draw_rectangle(tag.rect(), color = (0, 255, 0))
                #获取x,y,w,h
                rect = tag.rect()
                box = [rect[0], rect[1], rect[0] + rect[2], rect[1] + rect[3]]
                if calculate_IoU(box,last_frame_location) > 0.2:
                    flag_found_tag = True
                    x, y, w, h = rect[0], rect[1], rect[2], rect[3]
                    dx = (x - last_frame_x) / Ts
                    dy = (y - last_frame_y) / Ts
                    Z = np.array([x, y, w, h, dx, dy])
                    x_hat = Kalman_Filter(Z)
                    last_frame_x, last_frame_y = x, y
                    #img.draw_rectangle(last_frame_rect, color = (0, 0, 255))
                    img.draw_rectangle(rect, color = (255, 0, 0))
                    last_frame_rect = rect
                    last_frame_location = box
        if not flag_found_tag:
            predicted_rect = [
                int(x_hat[0]),
                int(x_hat[1]),
                int(x_hat[2]),
                int(x_hat[3])
            ]
            img.draw_rectangle(predicted_rect, color = (255, 255, 255))
            x,y,w,h = (x_hat[0] + (x_hat[4] * Ts)),(x_hat[1] + (x_hat[5] * Ts)),x_hat[2],x_hat[3]
            dx = (x - last_frame_x) / Ts
            dy = (y - last_frame_y) / Ts
            Z = np.array([x, y, w, h, dx, dy])
            x_hat = Kalman_Filter(Z)
            last_frame_x, last_frame_y = x, y
            last_frame_rect = [x,y,w,h]
            last_frame_location = [x,y,(x + w),(y + h)]